import os
import torch
from pathlib import Path

# 基础配置
BASE_DIR = Path(__file__).resolve().parent
DATA_DIR = BASE_DIR / "data"
MODEL_DIR = BASE_DIR / "models"

# 服务器配置
SERVER_HOST = "0.0.0.0"
SERVER_PORT = 8000
SERVER_URL = f"http://{SERVER_HOST}:{SERVER_PORT}"

# 联邦学习配置
FL_SERVER_HOST = os.getenv("FL_SERVER_HOST", "0.0.0.0")  # 可通过环境变量配置
FL_SERVER_PORT = int(os.getenv("FL_SERVER_PORT", "8080"))
FL_CLIENT_HOST = os.getenv("FL_CLIENT_HOST", "0.0.0.0")
FL_CLIENT_PORT = int(os.getenv("FL_CLIENT_PORT", "8081"))
MIN_CLIENTS = 2
ROUNDS = 3

# 模型配置
DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
BATCH_SIZE = 32
LEARNING_RATE = 0.001
EPOCHS = 5

# 创建必要的目录
os.makedirs(DATA_DIR, exist_ok=True)
os.makedirs(MODEL_DIR, exist_ok=True) 